import pandas as pd
from statsmodels.stats.outliers_influence import variance_inflation_factor
from statsmodels.tools.tools import add_constant
from sklearn.linear_model import LinearRegression
from common_import import *


def calculate_vif_and_explore_linear_relations(data, top_features):
    """
    计算给定数据中指定特征的VIF值，并探究VIF为无穷大的特征之间的线性关系。

    参数:
    - data: 包含所有特征的DataFrame
    - top_features: 要计算VIF的特征名列表

    返回:
    - vif_data: 包含特征名和VIF值的DataFrame
    - linear_relations: 线性关系的列表，其中包含符合条件的线性等式
    """
    # 选择前指定特征的数据
    df_top_features = data[top_features]

    # 增加一个常数列（为了计算VIF）
    X = add_constant(df_top_features)

    # 计算每个特征的VIF值
    vif_data = pd.DataFrame()
    vif_data["Feature"] = X.columns
    vif_data["VIF"] = [
        variance_inflation_factor(X.values, i) for i in range(X.shape[1])
    ]

    # 提取VIF为无穷大的特征
    inf_vif_features = vif_data[vif_data["VIF"] == float("inf")]["Feature"].tolist()

    # 如果没有VIF为无穷大的特征，直接返回
    if len(inf_vif_features) == 0:
        return vif_data, ["No features with infinite VIF found."]

    # 移除常数列
    inf_vif_features = [feature for feature in inf_vif_features if feature != "const"]

    # 存储线性等式
    linear_relations = []
    inf_vif_features = [feature for feature in inf_vif_features if "nBo" in feature]
    # 对每一个特征进行线性拟合
    for target_feature in inf_vif_features:
        # 剩余特征作为自变量
        other_features = [f for f in inf_vif_features if f != target_feature]
        X_others = data[other_features]
        y_target = data[target_feature]

        # 线性回归拟合
        linear_model = LinearRegression()
        linear_model.fit(X_others, y_target)
        r_squared = linear_model.score(X_others, y_target)

        # 如果R^2等于1，输出线性等式
        if r_squared > 0.99:
            intercept = linear_model.intercept_
            coefficients = linear_model.coef_
            # 构造线性等式字符串（去除系数为0的项）
            equation = f"{target_feature} = {intercept:.2f}"
            for coef, feature in zip(coefficients, other_features):
                if coef != 0:  # 排除系数为0的项
                    sign = "+" if coef > 0 else "-"
                    equation += f" {sign} {abs(coef):.2f} * {feature}"

            linear_relations.append(equation)

            # 构造线性等式字符串（去除系数为0的项）
            # flag = 0
            # equation = f"{target_feature} = {intercept:.2f}"
            # if abs(intercept) < 0.0001:
            #     equation = f"{target_feature} = "
            # for coef, feature in zip(coefficients, other_features):
            #     if (coef + 1) * (coef - 1) * coef > 0.001:
            #         flag = 1
            #     if abs(coef) > 0.00001:  # 排除系数为0的项
            #         sign = "+" if coef > 0 else "-"
            #         equation += f" {sign} {abs(coef):.2f} * {feature}"

            # if flag == 0:
            #     linear_relations.append(equation)

    return vif_data, linear_relations


# 读取包含前关注特征的文件
df_focus_feature = pd.read_csv("data/combined_importance.csv")
focus_feature = df_focus_feature["Feature"].tolist()
# 读取包含所有特征的数据文件
data = pd.read_csv("data/Molecular_Descriptor_training.csv")
data_focus = data[focus_feature]
# print(data_focus)
print(data_focus)

vif_data, linear_relations = calculate_vif_and_explore_linear_relations(
    data_focus, focus_feature
)
for i in linear_relations:
    print(i)
